package edu.washington.ling.syntric.proc;

import java.util.LinkedList;
import java.util.List;

import org.kohsuke.args4j.CmdLineException;
import org.kohsuke.args4j.CmdLineParser;

import edu.washington.ling.syntric.io.*;
import edu.washington.ling.syntric.struct.error.*;
import edu.washington.ling.syntric.struct.tree.*;
import edu.washington.ling.syntric.util.*;

public class MergeDSandPB {
	
	private static MergeDSandPBOptions opts = new MergeDSandPBOptions();
	private static PrintingExceptionHandler pHandler = new PrintingExceptionHandler();
	
	private static Token createMergedToken(Token dsToken, Token pbToken) {
		Token mergedToken = new Token(dsToken);
		for (String feature: pbToken.getFeatures()) {
			if (!mergedToken.hasFeature(feature) && !feature.equals("drel") && !feature.equals("dmrel")) {
				mergedToken.setFeature(feature, pbToken.getValue(feature));
			}
		}
		return mergedToken;
	}

	private static List<Token> mergeTokenLists(List<Token> dsTokens, List<Token> pbTokens) throws TreeMergeException {
		// create a list of nodes w/merged features containing (a) features from the DS node, 
		// and (b) features that are found only on the PB node.
		LinkedList<Token> dsTokenQ = new LinkedList<Token>(dsTokens);
		LinkedList<Token> pbTokenQ = new LinkedList<Token>(pbTokens);
		List<Token> mergedTokens = new LinkedList<Token>();
		while (dsTokenQ.size() > 0 && pbTokenQ.size() > 0) {
			Token dsToken = dsTokenQ.peekFirst();
			Token pbToken = pbTokenQ.peekFirst();
			if (!dsToken.isEmptyCategory() && !pbToken.isEmptyCategory()) {
				if (dsToken.getWord().equals(pbToken.getWord())) {
					dsTokenQ.removeFirst();
					pbTokenQ.removeFirst();
					mergedTokens.add(createMergedToken(dsToken, pbToken));
				}
				else {
					String dsString = MergeUtils.getSequenceString(dsTokens);
					String pbString = MergeUtils.getSequenceString(pbTokens);
					String errString = String.format("Words %s and %s do not match", dsToken.getWord(), pbToken.getWord());
					throw new TreeMergeException(errString, "DS", dsString, "PB", pbString);
				}
			}
			else if (dsToken.isDSEmptyCategory() && pbToken.isDSEmptyCategory()) {
				dsTokenQ.removeFirst();
				pbTokenQ.removeFirst();
				mergedTokens.add(createMergedToken(dsToken, pbToken));
			}
			else if (dsToken.isPSEmptyCategory() && pbToken.isPSEmptyCategory()) {
				if (dsToken.getValue("ectype").equals(pbToken.getValue("ectype"))) {
					dsTokenQ.removeFirst();
					pbTokenQ.removeFirst();
					mergedTokens.add(createMergedToken(dsToken, pbToken));
				}
				else {
					String dsString = MergeUtils.getSequenceString(dsTokens);
					String pbString = MergeUtils.getSequenceString(pbTokens);
					String errString = String.format("EC types %s and %s do not match", dsToken.getValue("ectype"), pbToken.getValue("ectype"));
					throw new TreeMergeException(errString, "DS", dsString, "PB", pbString);
				}
			}
			else if (!dsToken.isEmptyCategory() && pbToken.isPSEmptyCategory()) {
				pbTokenQ.removeFirst();
				mergedTokens.add(new Token(pbToken));
			}
			else {
				String dsString = MergeUtils.getSequenceString(dsTokens);
				String pbString = MergeUtils.getSequenceString(pbTokens);
				throw new TreeMergeException("Token merge error", "DS", dsString, "PB", pbString);
			}
		}
		if (dsTokenQ.size() > 0 || pbTokenQ.size() > 0) {
			String dsString = MergeUtils.getSequenceString(dsTokens);
			String pbString = MergeUtils.getSequenceString(pbTokens);
			throw new TreeMergeException("Token merge error", "DS", dsString, "PB", pbString);
		}
		return mergedTokens;
	}
	
	public static Tree mergeTrees(Tree ds, Tree pb) throws TreeMergeException {
		CollectingExceptionHandler cHandler = new CollectingExceptionHandler();
		Tree ds_pb = new Tree("DS+PB");
		List<Token> mergedTokens = mergeTokenLists(ds.getTokenList(), pb.getTokenList());
		ds_pb.addAllTokens(mergedTokens);
		// expand the drels and pbrels back into usable links. if there are any errors in expanding , say w/
		// the node reference labels (names) referred to by the pbrels (since we will be using names from DS) 
		// report those errors 
		try {
			TreeUtils.derriveDSLinksFromFeatures(ds_pb);
			TreeUtils.derrivePBLinksFromFeatures(ds_pb);
		}
		catch (UnreadableRepresentationException e){
			cHandler.handle(e);
			throw new TreeMergeException("Merged representation is unreadable", cHandler.getMessages());
		}
		return ds_pb;
	}

	private static void mergeTreesInInputFile(){
		int numExamples = 0;
		int numWithBoth = 0;
		int numMerged = 0;
		try {
			ExampleDocReader exfr = ExampleDocReader.newReader(opts.inputFilePath, "xml");
			ExampleDoc exf = exfr.read();
			ExampleDocWriter exfw = ExampleDocWriter.newWriter(opts.outputFilePath, "xml");
			TreeReader dsr = TreeReader.newReader(opts.dsTreeFormat);
			TreeReader pbr = TreeReader.newReader(opts.pbTreeFormat);
			TreeWriter trw = TreeWriter.newWriter(opts.dsTreeFormat);
			Example exr = null;
			while ((exr = exf.getNextExample()) != null) {
				numExamples++;
				if (exr.hasTreeRep("DS") && exr.hasTreeRep("PB")) {
					numWithBoth++;
					try {
						Tree ds = dsr.readTree(exr.getTreeRep("DS"));
						Tree pb = pbr.readTree(exr.getTreeRep("PB"));
						Tree ds_pb = mergeTrees(ds, pb);
						exr.addTreeRep(trw.writeTree(ds_pb));
						numMerged++;
					}
					catch (UnreadableRepresentationException e) {
						String errStr = String.format("Could not process representation for example '%s'", exr.getID());
						pHandler.handle(e, errStr);
					}
					catch (TreeMergeException e) {
						String errStr = String.format("Could not merge representations for example '%s'", exr.getID());
						pHandler.handle(e, errStr);
					}
				}
				exfw.writeExample(exr);
			}
			exfw.close();
			System.out.println("Merge DS and PB was successful.");
			System.out.println(String.format("%d/%d tree pairs were able to be merged", numMerged, numWithBoth));
			System.out.println(String.format("Processed %d total examples", numExamples));
		} 
		catch(Exception e) {
			pHandler.handle(e, "merge DS and PB failed.", System.out);
			e.printStackTrace(System.out);
			System.exit(-1);
		}
	}

	public static void main(String[] args) {
		CmdLineParser parser = new CmdLineParser(opts);
		try {
			parser.parseArgument(args);
			mergeTreesInInputFile();
		}
		catch(CmdLineException e) {
			System.out.println(e.getMessage());
			System.out.print("Usage: MergeDSandPB");
			parser.printSingleLineUsage(System.out);
			System.out.println();
			parser.printUsage(System.out);
			System.exit(-1);
		}
	}
}
